#!/usr/bin/python
# coding=utf-8

# from fancyimpute import MICE
import numpy as np
import scipy
import scipy.sparse
import scipy.sparse.linalg
import sys
import os

import random

from .re_utils import *
from .extract_controller_point import get_model_all_measure_value

# A Reshaper provide multiple methods for synthesizing body mesh
class Reshaper:
  def __init__(self, label="female"):
    self.label_ = label
    self.facets = np.load(open(os.path.join(MODEL_DIR, "facets.npy"), "rb"))
    self.normals = np.load(open(os.path.join(MODEL_DIR, "normals.npy"), "rb"))
    # self.mask = np.load(open(os.path.join(MODEL_DIR, "mask.npy"), "rb"))
    self.rfemask = np.load(open(os.path.join(MODEL_DIR, "%s_rfemask.npy"%label), "rb"))
    self.rfemat = np.load(open(os.path.join(MODEL_DIR, "%s_rfemat.npy"%label), "rb"))
    self.m2d = np.load(open(os.path.join(MODEL_DIR, "%s_m2d.npy"%label), "rb"))
    self.d_basis = np.load(open(os.path.join(MODEL_DIR, "%s_d_basis.npy"%label), "rb"))
    self.t_measure = np.load(open(os.path.join(MODEL_DIR, "%s_t_measure.npy"%label), "rb"))
    self.mean_measure = np.load(open(os.path.join(MODEL_DIR, "%s_mean_measure.npy"%label), "rb"))
    self.mean_deform = np.load(open(os.path.join(MODEL_DIR, "%s_mean_deform.npy"%label), "rb"))
    self.mean_vertex = np.load(open(os.path.join(MODEL_DIR, "%s_mean_vertex.npy"%label), "rb"))
    self.std_measure = np.load(open(os.path.join(MODEL_DIR, "%s_std_measure.npy"%label), "rb"))
    self.std_deform = np.load(open(os.path.join(MODEL_DIR, "%s_std_deform.npy"%label), "rb"))
    # self.cp = np.load(open(os.path.join(MODEL_DIR, "cp.npy"), "rb"))

    loader = np.load(os.path.join(MODEL_DIR, "%s_d2v.npz"%label))
    # 以坐标的形式构造稀疏函数
    self.d2v = scipy.sparse.coo_matrix((loader['data'], (loader['row'], loader['col'])),shape=loader['shape'])
    self.lu = scipy.sparse.linalg.splu(self.d2v.transpose().dot(self.d2v).tocsc())
    # self.local_mat = []
    # tmp = np.load(open(os.path.join(MODEL_DIR, "%s_local.npy"%label), "rb"))
    # for i in range(0, len(tmp)):
    #   self.local_mat.append(np.array([c for c in tmp[i]]))

  def mapping(self, weight, flag=0):
    if flag==0:
      return self.mapping_global(weight)
    elif flag==1:
      return self.mapping_mask(weight)
    else:
      return self.mapping_rfemat(weight)

  # global mapping using t_measure
  def mapping_global(self, weight):
    weight = np.array(weight).reshape(M_NUM, 1)
    weight = self.m2d.dot(weight)
    d = np.matmul(self.d_basis, weight)
    d.shape = (F_NUM, 9)
    d *= self.std_deform
    d += self.mean_deform
    d.shape = (F_NUM * 9, 1)
    [v, n, f] = self.d_synthesize(d)
    return [v, n, f]

  # local mapping using measure + mask
  def mapping_mask(self, weight):
    # weight = np.array(weight).reshape(M_NUM, 1)
    # weight *= self.std_measure
    # weight += self.mean_measure
    d = []
    for i in range(0, F_NUM):
      mask = np.array(self.mask[:, i]).reshape(M_NUM, 1)
      alpha = np.array(weight[mask])
      alpha.shape = (alpha.size, 1)
      s = self.local_mat[i].dot(alpha)
      d += [a for a in s.flat]
    d = np.array(d).reshape(F_NUM * 9, 1)
    [v, n, f] = self.d_synthesize(d)
    return [v, n, f]

  # local mapping using measure + rfe_mat
  def mapping_rfemat(self, weight):
    # 测量尺寸反归一化
    weight = np.array(weight).reshape(M_NUM, 1)
    # weight *= self.std_measure
    # weight += self.mean_measure
    d = []
    for i in range(0, F_NUM):
      # 选掩膜
      mask = np.array(self.rfemask[:, i]).reshape(M_NUM, 1)
      # 筛选测量参数
      alpha = np.array(weight[mask])
      alpha.shape = (alpha.size, 1)
      # 求该面的变形矩阵
      s = self.rfemat[i].dot(alpha)
      d += [a for a in s.flat]
    # 该样本的变形集
    d = np.array(d).reshape(F_NUM * 9, 1)
    [v, n, f] = self.d_synthesize(d)
    return [v, n, f]

  # synthesize a body by deform-based, given deform, output vertex
  def d_synthesize(self, deform):
    d = np.array(deform.flat).reshape(deform.size, 1)
    Atd = self.d2v.transpose().dot(d)
    x = self.lu.solve(Atd)
    x = x[:V_NUM * 3]
    # move to center
    x.shape = (V_NUM, 3)
    x_mean = np.mean(x, axis=0)
    x -= x_mean
    return [x, -self.normals, self.facets - 1]

  # given flag, value, predict completed measures
  def test(self, flag, data):
    if (flag == 1).sum() == self.data.m_num:
      return data
    else:
      solver = MICE()
      return self.imputate(flag, data, solver)

  # using imputation for missing data
  def get_predict(self, flag, in_data):
    output = in_data.copy()
    output.shape = (M_NUM, 1)
    output[~flag] = np.nan
    solver = MICE()
    tmp = self.t_measure.copy()
    tmp = np.column_stack((tmp, output)).transpose()
    tmp = solver.complete(tmp)
    output = np.array(tmp[-1, :]).reshape(M_NUM, 1)
    return output


def reshape_contrast(body, measure_val, type_s, choice_num, label='male'):
  sample_list = [i for i in range(len(measure_val[0]))]
  sample_list = random.sample(sample_list, choice_num)
  measure_val = measure_val[:,sample_list]

  row_name = ['height', 'shoulder_width', 'arm_length', 'leg_length', 'chest_length', 'chest_width', 'chest_thickness', 'waist_length', 'waist_width', 'waist_thickness', 'hips_length', 'hips_width', 'hips_thickness', 'right-leg_length', 'right-leg_width', 'right-leg_thickness']
  err_arr = []
  for i in range(len(measure_val[0])):
    [v, n, f] = body.mapping(measure_val[:,i], type_s)
    save_obj(os.path.join(OUTPUT_DIR, str(i) + "_global.obj"), v, f+1)
    reshape_dir = get_model_all_measure_value(os.path.join(OUTPUT_DIR, str(i) + "_global.obj"))
    reshape_arr = []
    for j in range(len(row_name)):
      reshape_arr.append(reshape_dir[row_name[j]])
    reshape_arr = np.array(reshape_arr).reshape(M_NUM,)
    err_arr.append(abs(measure_val[:,i] - reshape_arr))
  return np.array(err_arr).mean(0)


# 计算误差项是否合要求
def cal_err_whether_in_limit(err_arr, limit):
  status = abs(err_arr) > limit
  # print('status:', status)
  return status.any()

# 获取惩罚项后的输入
def get_penalty_input(last_err_arr, err_arr, p, measure):
  if (len(last_err_arr) > 0):
    # print(err_arr)
    # print('penalty before:', measure)
    optimize_res = abs(err_arr) - abs(last_err_arr) > 0
    # print('penalty status:', optimize_res)
    # print('penalty after:', measure)
    for i in range(len(optimize_res)):
      if optimize_res[i]:
        measure[i] = measure[i] - p * err_arr[i]
    # print('after:', measure)
  return measure

# 获取模型重建的误差
def get_reshape_error(body, target, measure_val, label, index=0):
  row_name = ['height', 'shoulder_width', 'arm_length', 'leg_length', 'chest_length', 'chest_width', 'chest_thickness', 'waist_length', 'waist_width', 'waist_thickness', 'hips_length', 'hips_width', 'hips_thickness', 'right-leg_thickness', 'head_length', 'body_length']
  [v, n, f] = body.mapping(measure_val, 2)
  save_obj(os.path.join(OUTPUT_DIR, str(index) + label + "_optimize.obj"), v, f+1)
  modelPath = os.path.join(OUTPUT_DIR, str(index) + label + "_optimize.obj")
  reshape_dir = get_model_all_measure_value(modelPath)
  reshape_arr = []
  for j in range(len(row_name)):
    reshape_arr.append(reshape_dir[row_name[j]])
  return (np.array(target) - np.array(reshape_arr)), modelPath

# 优化方式的模型重建
def with_optimize_reshape(body, measure_val, label='male', limit=5, isOptimize=True):
  err_arr = []
  last_err_arr = []
  target = measure_val
  w = 0.6
  p = 0.15
  index = 0

  err_arr, modelPath = get_reshape_error(body, target, measure_val, label, index=index)

  while cal_err_whether_in_limit(err_arr, limit) and isOptimize:
    print(err_arr)
    # 修正项加入
    measure_val = measure_val + w * err_arr
    measure_val = get_penalty_input(last_err_arr, err_arr, p, measure_val)

    last_err_arr = err_arr
    index = index + 1
    err_arr, modelPath = get_reshape_error(body, target, measure_val, label, index=index)
    # print(err_arr)
  return err_arr, modelPath


# 测试数据优化对比
def optimize_contrast(body, measure_set, choice_num, label, limit):
  err_set = []

  sample_list = [i for i in range(len(measure_set[0]))]
  sample_list = random.sample(sample_list, choice_num)
  measure_set = measure_set[:,sample_list]

  for i in range(len(measure_set[0])):
    err_set.append(with_optimize_reshape(body, measure_set[:,i], label='male', limit=5))
  return np.array(err_set).mean(0)


if __name__ == "__main__":
  DATA_DIR = "new_data"
  label = "female"
  body = Reshaper(label)
  
  measure = np.load(open(os.path.join(DATA_DIR, label + "_measure.npy"),"rb"))
  # print(measure[:, 2])
  # print(reshape_contrast(body, measure, 0, 100, label))
  print(with_optimize_reshape(body, measure[:, 2], label, 1))
  # print(optimize_contrast(body, measure, 100, label, 5))

  # 局部映射测试
  # measure = [1579,348,411,641,955,305,256,805,290,217,1059,381,271,644,190,210]
  # int_float_map = map(float, measure)
  # measure = list(int_float_map)
  # measure = np.array(measure).reshape(M_NUM, 1)
  # t_measure = measure - mean_measure
  # t_measure /= std_measure
  # # [v, n, f] = body.mapping(t_measure, 2)
  # [v, n, f] = body.mapping(t_measure, 0)
  
  # save_obj(os.path.join(OUTPUT_DIR, "test.obj"), v, f+1)


  # for i in range(100):
  #   [v, n, f] = body.mapping(t_measure[:,i], 2)
  #   save_obj(os.path.join(MODEL_DIR, "test.obj"), v, f+1)


